#coding:utf8
import numpy as np
np.random.seed(123)
import os
#后面只使用keras.model搭建一个简单的全连接网络模型，不用tf.keras中的特性，在此直接用import keras也可以
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from keras.utils import np_utils
from tensorflow.keras.optimizers import RMSprop
from MNISTModel import create_model
from DataGenerator import *

if __name__ == '__main__':

    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    physical_devices = tf.config.experimental.list_physical_devices('GPU')
    assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

    # 直接使用keras.Sequential()搭建全连接网络模型
    model=create_model()
    x_train, y_train=getTrainData()
    model.fit(x_train,y_train,epochs=5,batch_size=32)
    model.save_weights("./resources/MNIST_model_swlw/")
    print("模型已保存")